// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
//

#include <kllm/kai/chat_completions.h>
#include <kllm/utility/utils.h>

namespace kllm {

    void KaiChatCompletions::chat(const KaiRequest &req, KaiResponse &res) {
        res.set_model_alias(_context->params.model_alias);
        PromptsValue value;
        value.set_string_value(req.chat_completion_request().message());
        KaiPrompts prompts;
        *prompts.mutable_values()->Add() = std::move(value);
        std::vector<server_task> tasks = _context->create_tasks_inference(req, SERVER_TASK_INF_TYPE_COMPLETION, prompts);
        _context->queue_results.add_waiting_tasks(tasks);
        _context->queue_tasks.post(tasks);

        const auto task_ids = server_task::get_list_id(tasks);

        const auto completion_id = gen_chatcmplid();

        _context->receive_cmpl_results(task_ids, [&](std::vector<TaskResult> & results) {
            if (results.size() == 1) {
                // single result
                auto r = results[0].completions;
                r.set_completion_id(completion_id);
                *res.mutable_completions()->Add() = std::move(r);
            } else {
                // multiple results (multitask)
                for (const auto & r : results) {
                    auto rr = r.completions;
                    rr.set_completion_id(completion_id);
                    *res.mutable_completions()->Add() = std::move(rr);
                }
            }
        }, [&](const turbo::Status & status) {
            res.mutable_status()->set_code(status.raw_code());
            res.mutable_status()->set_errmsg(std::string(status.message()));
        });

        _context->queue_results.remove_waiting_task_ids(task_ids);
    }

    void KaiChatCompletions::chat_stream(const KaiRequest &req, KaiResponse &res,
                     const std::function<bool(const CompletionsResult &)> &func,
                     const std::function<void()> &on_complete) {
        res.set_model_alias(_context->params.model_alias);
        PromptsValue value;
        value.set_string_value(req.chat_completion_request().message());
        KaiPrompts prompts;
        *prompts.mutable_values()->Add() = std::move(value);
        std::vector<server_task> tasks = _context->create_tasks_inference(req, SERVER_TASK_INF_TYPE_COMPLETION, prompts);
        _context->queue_results.add_waiting_tasks(tasks);
        _context->queue_tasks.post(tasks);
        const auto task_ids = server_task::get_list_id(tasks);
        const auto completion_id = gen_chatcmplid();
        auto context_call_back = [&](const TaskResult & result) -> bool {
            auto c = result.completions;
            c.set_completion_id(completion_id);
            if(c.model().empty()) {
                c.set_model(std::string(DEFAULT_OAICOMPAT_MODEL));
            }
            *res.mutable_completions()->Add() = c;
            auto r = func(c);
            LOG(INFO)<<"call result: "<<r;
            return r;
        };
        auto context_err_back = [&](const turbo::Status & status) {
            res.mutable_status()->set_code(status.raw_code());
            res.mutable_status()->set_errmsg(std::string(status.message()));
            LOG(INFO)<<"on error";
            on_complete();
        };

        this->_context->receive_cmpl_results_stream(task_ids, context_call_back, context_err_back);
        _context->queue_results.remove_waiting_task_ids(task_ids);
        LOG(INFO)<<"call done";
        on_complete();

    }

}  // namespace kllm
